import gym
from engine.account import Account
import numpy as np

class PortfolioEnv(gym.Env):
    def __init__(self, df):
        self.df = df
        self.dates = list(df.index.unique())
        self.acc = Account()
        super(PortfolioEnv, self).__init__()

    def softmax_normalization(self, actions):
        numerator = np.exp(actions)
        denominator = np.sum(np.exp(actions))
        softmax_output = numerator / denominator
        return softmax_output

    def step(self, actions):
        #actions变成weights,根据return与新权重更新市值。
        #新市值=reward，state是数据
        self.terminal = self.day >= len(self.df.index.unique()) - 1
        if self.terminal:
            print('组合终值：',self.acc.get_total_mv())
            print('done!')
            return self.state, self.reward, self.terminal, {}
        self.day += 1
        date = self.dates[self.day]
        df_bar = self.df.loc[date]

        self.acc.update_bar(date, df_bar)
        weights = self.softmax_normalization(actions)
        if weights is not None:
            self.acc.adjust_weights(date, weights)

        total_mv = self.acc.get_total_mv()
        self.reward = total_mv

        features = []
        self.state = df_bar[features]

        return self.state, self.reward, self.terminal, {}

    def reset(self):
        self.day = 0

    def render(self, mode='human'): #渲染引擎
        pass

if __name__ == '__main__':
    from logic.globalpy import D
    from engine.data.data_handler import DataHandler

    fields, names = DataHandler().get_kbar_fields_names()

    fields.append('Ref($close,-1)/$close -1')
    names.append('label_value')

    fields.append('QCut($label_value,10)')
    names.append('label')

    # D = Dataloader(path='../../config/indexes')
    df_all = D.load(['000300.SH', '000905.SH', '399006.SZ'], start_time='20100101', fields=fields, names=names)
    print(df_all.head())
    print(df_all.tail(30))

    env = PortfolioEnv(df=df_all)
    print(env)
